{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Finetune and deploy Cohere Classification Models from AWS Marketplace\n",
    "\n",
    "This sample notebook shows you how to finetune and deploy classification models using Amazon SageMaker.\n",
    "\n",
    "> **Note**: This is a reference notebook and it cannot run unless you make changes suggested in the notebook.\n",
    "\n",
    "> The classifications-finetuning model package supports SageMaker Realtime Inference but not SageMaker Batch Transform.\n",
    "\n",
    "## Pre-requisites:\n",
    "1. **Note**: This notebook contains elements which render correctly in Jupyter interface. Open this notebook from an Amazon SageMaker Notebook Instance or Amazon SageMaker Studio.\n",
    "1. Ensure that IAM role used has **AmazonSageMakerFullAccess**\n",
    "1. To deploy this ML model successfully, ensure that:\n",
    "    1. Either your IAM role has these three permissions and you have authority to make AWS Marketplace subscriptions in the AWS account used: \n",
    "        1. **aws-marketplace:ViewSubscriptions**\n",
    "        1. **aws-marketplace:Unsubscribe**\n",
    "        1. **aws-marketplace:Subscribe**  \n",
    "    2. or your AWS account has a subscription to [cohere-classification-finetuning-multilingual](https://aws.amazon.com/marketplace/pp/prodview-hukrkor45t7bw). If so, skip step: [Subscribe to the model package](#1.-Subscribe-to-the-model-package)\n",
    "\n",
    "## Contents:\n",
    "1. [Subscribe to the model package](#1.-Subscribe-to-the-model-package)\n",
    "2. [Finetune Classification Models](#2.-Finetune-Classification-Models)\n",
    "   1. [Upload training data](#A.-Upload-training-data)\n",
    "   2. [Finetune models on uploaded data](#B.-Finetune-models-on-uploaded-data)\n",
    "3. [Create an endpoint for inference with multiple models](#3.-Create-an-endpoint-for-inference-with-multiple-models)\n",
    "   1. [Create an endpoint]()\n",
    "   2. [Perform real-time inference]()\n",
    "4. [Clean-up](#4.-Clean-up)\n",
    "    1. [Delete the endpoint](#A.-Delete-the-endpoint)\n",
    "    2. [Unsubscribe to the listing (optional)](#B.-Unsubscribe-to-the-listing-(optional))\n",
    "    \n",
    "\n",
    "## Usage instructions\n",
    "You can run this notebook one cell at a time (By using Shift+Enter for running a cell)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Subscribe to the algorithm"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To subscribe to the model package:\n",
    "1. Open the algorithm listing page [cohere-classification-finetuning-multilingual](https://aws.amazon.com/marketplace/pp/prodview-hukrkor45t7bw)\n",
    "1. On the AWS Marketplace listing, click on the **Continue to subscribe** button.\n",
    "1. On the **Subscribe to this software** page, review and click on **\"Accept Offer\"** if you and your organization agrees with EULA, pricing, and support terms. \n",
    "1. Once you click on **Continue to configuration button** and then choose a **region**, you will see a **Product Arn** displayed. This is the algorithm ARN that you need to specify while creating a deployable model using Boto3. Copy the ARN corresponding to your region and specify the same in the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --upgrade cohere-aws\n",
    "# if you upgrade the package, you need to restart the kernel\n",
    "\n",
    "from cohere_aws import Client\n",
    "import boto3\n",
    "import sagemaker as sage\n",
    "from sagemaker.s3 import S3Uploader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "region = boto3.Session().region_name\n",
    "\n",
    "cohere_package = \"classification-finetuning-mult-244c14a2987a3828b545545ce8884346\"\n",
    "\n",
    "algorithm_map = {\n",
    "    \"us-east-1\": f\"arn:aws:sagemaker:us-east-1:865070037744:algorithm/{cohere_package}\",\n",
    "    \"us-east-2\": f\"arn:aws:sagemaker:us-east-2:057799348421:algorithm/{cohere_package}\",\n",
    "    \"us-west-1\": f\"arn:aws:sagemaker:us-west-1:382657785993:algorithm/{cohere_package}\",\n",
    "    \"us-west-2\": f\"arn:aws:sagemaker:us-west-2:594846645681:algorithm/{cohere_package}\",\n",
    "    \"ca-central-1\": f\"arn:aws:sagemaker:ca-central-1:470592106596:algorithm/{cohere_package}\",\n",
    "    \"eu-central-1\": f\"arn:aws:sagemaker:eu-central-1:446921602837:algorithm/{cohere_package}\",\n",
    "    \"eu-west-1\": f\"arn:aws:sagemaker:eu-west-1:985815980388:algorithm/{cohere_package}\",\n",
    "    \"eu-west-2\": f\"arn:aws:sagemaker:eu-west-2:856760150666:algorithm/{cohere_package}\",\n",
    "    \"eu-west-3\": f\"arn:aws:sagemaker:eu-west-3:843114510376:algorithm/{cohere_package}\",\n",
    "    \"eu-north-1\": f\"arn:aws:sagemaker:eu-north-1:136758871317:algorithm/{cohere_package}\",\n",
    "    \"ap-southeast-1\": f\"arn:aws:sagemaker:ap-southeast-1:192199979996:algorithm/{cohere_package}\",\n",
    "    \"ap-southeast-2\": f\"arn:aws:sagemaker:ap-southeast-2:666831318237:algorithm/{cohere_package}\",\n",
    "    \"ap-northeast-2\": f\"arn:aws:sagemaker:ap-northeast-2:745090734665:algorithm/{cohere_package}\",\n",
    "    \"ap-northeast-1\": f\"arn:aws:sagemaker:ap-northeast-1:977537786026:algorithm/{cohere_package}\",\n",
    "    \"ap-south-1\": f\"arn:aws:sagemaker:ap-south-1:077584701553:algorithm/{cohere_package}\",\n",
    "    \"sa-east-1\": f\"arn:aws:sagemaker:sa-east-1:270155090741:algorithm/{cohere_package}\",\n",
    "}\n",
    "if region not in algorithm_map.keys():\n",
    "    raise Exception(f\"Current boto3 session region {region} is not supported.\")\n",
    "\n",
    "arn = algorithm_map[region]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Finetune classification models"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Upload training data\n",
    "\n",
    "Choose a directory on S3 to store the training data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_data_dir = \"s3/...\"  # where to upload the data"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Upload sample training data to S3:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sess = sage.Session()\n",
    "train_dataset1 = S3Uploader.upload(\"../examples/sample_sentiment_classification_data.jsonl\", s3_data_dir.rstrip(\"/\"), sagemaker_session=sess)\n",
    "train_dataset2 = S3Uploader.upload(\"../examples/sample_multilabel_classification_data.jsonl\", s3_data_dir.rstrip(\"/\"), sagemaker_session=sess)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### B. Finetune models on uploaded data\n",
    "\n",
    "Specify a directory on S3 where finetuned models should be stored:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_models_dir = \"s3/...\"  # where the models will be saved "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create Cohere client:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "co = Client(region_name=region)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create fine-tuning jobs for both uploaded datasets:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "co.create_finetune(arn=arn, name=\"model1\", train_data=train_dataset1, s3_models_dir=s3_models_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "co.create_finetune(arn=arn, name=\"model2\", train_data=train_dataset2, s3_models_dir=s3_models_dir)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Create an endpoint for inference with multiple models"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Create an endpoint\n",
    "\n",
    "The Cohere AWS SDK provides a built-in method for creating an endpoint for inference. This will automatically deploy all models you finetuned earlier.\n",
    "\n",
    "> **Note**: This is equivalent to creating and deploying a `ModelPackage` in SageMaker's SDK."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "co.create_endpoint(arn=arn, endpoint_name=\"cohere-multilingual-classify\", s3_models_dir=s3_models_dir, recreate=True)\n",
    "\n",
    "# # If the endpoint is already created, you just need to connect to it\n",
    "# co.connect_to_endpoint(endpoint_name=\"cohere-multilingual-classify\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### B. Perform real-time inference\n",
    "\n",
    "Now, you can access all models deployed on the endpoint for inference:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(co.classify([\"it works!\", \"cela fonctionne!\"], \"model1\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(co.classify([\"Vamos a la pizzeria\"], \"model2\"))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Clean-up"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Delete the endpoint\n",
    "\n",
    "After you've successfully performed inference, you can delete the deployed endpoint to avoid being charged continuously. This can also be done via the Cohere AWS SDK:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# co.delete_endpoint()\n",
    "# co.close()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### B. Unsubscribe to the listing (optional)\n",
    "\n",
    "If you would like to unsubscribe to the model package, follow these steps. Before you cancel the subscription, ensure that you do not have any [deployable model](https://console.aws.amazon.com/sagemaker/home#/models) created from the model package or using the algorithm. Note - You can find this information by looking at the container name associated with the model. \n",
    "\n",
    "**Steps to unsubscribe to product from AWS Marketplace**:\n",
    "1. Navigate to __Machine Learning__ tab on [__Your Software subscriptions page__](https://aws.amazon.com/marketplace/ai/library?productType=ml&ref_=mlmp_gitdemo_indust)\n",
    "2. Locate the listing that you want to cancel the subscription for, and then choose __Cancel Subscription__  to cancel the subscription.\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "emb",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
